"""
@Filename       : uikg_construct.py
@Create Time    : 2022/3/8 23:09
@Author         : Rylynn
@Description    : 

"""
import json

import dgl
import networkx as nx
import os.path

from util.preprocess import load_vocab_dict, load_content


def load_uikg(root_path, dataset):
    print('Loading UIKG...')
    kg = json.load(open(os.path.join(root_path, dataset, 'content.json').format(dataset)))
    vocab_dict = load_vocab_dict(root_path, dataset)
    content_dict = load_content(os.path.join(root_path, dataset))
    entity_dict = {}
    max_eid = 0
    with open(os.path.join(root_path, dataset, 'entities.tsv')) as f:
        for line in f.readlines():
            eid, entity = line.strip().split('\t')
            entity_dict[entity] = int(eid)

    uikg_dict = {}
    for user, uid in vocab_dict.items():
        uikg_dict[uid] = nx.Graph()


    with open(os.path.join(root_path, dataset, 'cascade.txt')) as f:
        for line in f.readlines():
            content_id = line.strip().split()[0]
            if not content_dict.get(content_id) or not kg.get(content_id):
                continue

            chunks = line.strip().split()[1:]
            triples = kg[content_id]
            edges_list = []

            for h, r, t in triples:
                if entity_dict.get(h) and entity_dict.get(t):
                    edges_list.append((entity_dict[h], entity_dict[t]))
                    edges_list.append((entity_dict[t], entity_dict[h]))

            for chunk in chunks:
                user, timestamp = chunk.split(',')
                user, timestamp = user, eval(timestamp)
                if vocab_dict.get(user):
                    user_id = vocab_dict[user]
                    uikg_dict[user_id].add_edges_from(edges_list)

    mean_entity_number = sum(map(lambda x: x[1].number_of_nodes(), uikg_dict.items())) / len(uikg_dict)
    mean_link_number = sum(map(lambda x: x[1].number_of_edges(), uikg_dict.items())) / len(uikg_dict)

    empty_graph_uid = []
    for user_id, uikg in uikg_dict.items():
        if uikg.number_of_nodes() == 0:
            empty_graph_uid.append(user_id)
            continue
        origin_nid = list(uikg.nodes)
        nid_relabel_dict = {nid: idx for idx, nid in enumerate(origin_nid)}
        relabeled_uikg = nx.relabel_nodes(uikg, nid_relabel_dict)
        relabeled_uikg = dgl.from_networkx(relabeled_uikg).to('cuda')
        uikg_dict[user_id] = (relabeled_uikg, origin_nid)

    for id in empty_graph_uid:
        uikg_dict.pop(id)

    print('Mean entity number: ', mean_entity_number)
    print('Mean link number: ', mean_link_number)
    print('Loading FINISH...')
    # for uid, g in uikg_dict.items():
    #     g.add_edge(48006, 48006)
    #     uikg_dict[uid] = g.to('cuda')

    return uikg_dict
